import matplotlib

matplotlib.use('TkAgg')  # 或者 'Qt5Agg'：设置Matplotlib使用TkAgg或Qt5Agg作为图形后端，适用于图形界面显示。
import matplotlib.pyplot as plt  # 导入Matplotlib的pyplot模块，用于绘制图表。
import torch  # 导入PyTorch库，用于深度学习模型的构建与训练。
import torch.nn as nn  # 导入PyTorch的神经网络模块。
import numpy as np  # 导入NumPy库，用于数值计算。
import sys  # 导入sys库，提供对Python运行时环境的访问。

sys.path.append("..")  # 将父目录添加到Python模块搜索路径中，便于导入自定义模块。
import d2lzh_pytorch as d2l  # 导入d2lzh_pytorch自定义模块，提供一些深度学习相关的辅助函数。
import pandas as pd  # 导入Pandas库，用于数据处理和操作。

# 设置默认数据类型为 float32 (相当于原来的 FloatTensor)
torch.set_default_dtype(torch.float32)  # 设置PyTorch的默认数据类型为float32。

# 读取训练数据集与测试数据集
train_data = pd.read_csv('dataset/train.csv')  # 读取训练数据集，返回一个DataFrame。
test_data = pd.read_csv('dataset/test.csv')  # 读取测试数据集，返回一个DataFrame。

# 合并训练数据和测试数据的特征
all_features = pd.concat((train_data.iloc[:, 1:-1], test_data.iloc[:, 1:]))
# 将训练数据集的特征（排除第1列ID和最后一列标签）和测试数据集的特征（排除第1列ID）合并。

# 识别出所有数值型特征,找到标签列
numeric_features = all_features.dtypes[all_features.dtypes != 'object'].index  # 识别数据中所有数值型特征。
# print(numeric_features)  # 打印所有数值型特征的列名。

# 数据标准化为均值为0，标准差为1。
all_features[numeric_features] = all_features[numeric_features].apply(lambda x: (x - x.mean()) / (x.std()))
# 对所有数值型特征进行标准化，使每个特征的均值为0，标准差为1。

# 标准化后，每个特征的均值变为0，所以可以直接用0来替换缺失值
all_features = all_features.fillna(0)  # 将所有缺失值替换为0。

# 将离散值转换成指示特征
all_features = pd.get_dummies(all_features, dummy_na=True)
# 将所有离散特征转换为指示特征（one-hot encoding）。dummy_na=True表示将缺失值也作为一种特征。

# 将True和False转成0、1
all_features = all_features * 1  # 将所有的True转为1，False转为0。

# 将数据转换成NumPy格式的数据，并转成 NDArray 方便后面的训练
n_train = train_data.shape[0]  # 获取训练数据集的样本数量。
train_features = torch.tensor(all_features[:n_train].values, dtype=torch.float)
# 提取训练集的特征并转换成PyTorch张量（tensor）。
test_features = torch.tensor(all_features[n_train:].values, dtype=torch.float)
# 提取测试集的特征并转换成PyTorch张量（tensor）。
train_labels = torch.tensor(train_data.SalePrice.values, dtype=torch.float).view(-1, 1)
# 提取训练集的标签（房价），并将其转换为列向量形式。

# 基本线性回归模型、平方损失函数
loss = torch.nn.MSELoss()  # 定义平方损失函数（Mean Squared Error Loss），用于回归问题。


def get_net(feature_num):
    # 定义一个基本的线性回归网络
    net = nn.Linear(feature_num, 1)  # 使用PyTorch的线性层（全连接层），输入特征数为feature_num，输出1个值（预测房价）。
    for param in net.parameters():
        nn.init.normal_(param, mean=0, std=0.01)  # 初始化网络参数，使用正态分布进行初始化，均值为0，标准差为0.01。
    return net  # 返回模型。


# 对数均方根误差
def log_rmse(net, features, labels):
    with torch.no_grad():  # 在计算误差时不需要计算梯度。
        # 将小于1的值设成1，使得取对数时数值更稳定
        clipped_preds = torch.max(net(features), torch.tensor(1.0))
        # 计算对数均方根误差（log RMSE），先计算预测值的对数和标签值的对数，然后计算RMSE。
        rmse = torch.sqrt(2 * loss(clipped_preds.log(), labels.log()).mean())  # 计算对数均方根误差。
    return rmse.item()  # 返回误差值。


# 使用Adam优化算法
def train(net, train_features, train_labels, test_features, test_labels, num_epochs, learning_rate, weight_decay,
          batch_size):
    train_ls, test_ls = [], []  # 用于存储训练集和测试集的误差值。
    dataset = torch.utils.data.TensorDataset(train_features, train_labels)  # 将训练数据和标签打包成一个数据集。
    train_iter = torch.utils.data.DataLoader(dataset, batch_size, shuffle=True)  # 创建一个DataLoader，按batch_size批次加载训练数据。
    # 这里使用Adam优化算法，优化模型参数
    optimizer = torch.optim.Adam(params=net.parameters(), lr=learning_rate, weight_decay=weight_decay)
    net = net.float()  # 将模型转换为浮点数类型。

    for epoch in range(num_epochs):  # 训练num_epochs轮。
        for X, y in train_iter:  # 遍历每个batch的训练数据。
            l = loss(net(X.float()), y.float())  # 计算当前模型的损失。
            optimizer.zero_grad()  # 清空优化器中的梯度。
            l.backward()  # 反向传播计算梯度。
            optimizer.step()  # 更新模型参数。
        train_ls.append(log_rmse(net, train_features, train_labels))  # 计算并记录训练集的误差。
        if test_labels is not None:  # 如果有测试集标签，计算并记录测试集的误差。
            test_ls.append(log_rmse(net, test_features, test_labels))
    return train_ls, test_ls  # 返回训练集和测试集的误差。


# K折交叉验证,用于评估模型表现，确保模型的泛化能力
def get_k_fold_data(k, i, X, y):
    # 返回第i折交叉验证时所需要的训练和验证数据
    assert k > 1  # 确保k大于1，因为K折交叉验证至少需要2折
    fold_size = X.shape[0] // k  # 计算每一折的样本数量
    X_train, y_train = None, None  # 初始化训练集和标签
    for j in range(k):  # 遍历每一折
        idx = slice(j * fold_size, (j + 1) * fold_size)  # 获取当前折的索引范围
        X_part, y_part = X[idx, :], y[idx]  # 获取当前折的数据和标签
        if j == i:  # 如果是第i折，则将其作为验证集
            X_valid, y_valid = X_part, y_part
        elif X_train is None:  # 如果是第一次定义训练集，则直接赋值
            X_train, y_train = X_part, y_part
        else:  # 将剩余的折拼接成训练集
            X_train = torch.cat((X_train, X_part), dim=0)
            y_train = torch.cat((y_train, y_part), dim=0)
    return X_train, y_train, X_valid, y_valid  # 返回训练集和验证集


# 在K折交叉验证中我们训练K次并返回训练和验证的平均误差。
def k_fold(k, X_train, y_train, num_epochs, learning_rate, weight_decay, batch_size):
    train_l_sum, valid_l_sum = 0, 0  # 初始化训练误差和验证误差的累积
    for i in range(k):  # 执行k次交叉验证
        data = get_k_fold_data(k, i, X_train, y_train)  # 获取第i折的训练和验证数据
        net = get_net(X_train.shape[1])  # 创建网络
        train_ls, valid_ls = train(net, *data, num_epochs, learning_rate, weight_decay, batch_size)  # 训练模型并获取误差
        train_l_sum += train_ls[-1]  # 累加训练误差
        valid_l_sum += valid_ls[-1]  # 累加验证误差
        if i == 0:  # 如果是第一次，绘制训练和验证误差曲线
            d2l.semilogy(range(1, num_epochs + 1), train_ls, 'epochs', 'rmse', range(1, num_epochs + 1), valid_ls,['train', 'valid'], plot=True)
        print('fold %d, train rmse %f, valid rmse %f' % (i, train_ls[-1], valid_ls[-1]))  # 打印当前折的训练和验证误差
    return train_l_sum / k, valid_l_sum / k  # 返回训练和验证的平均误差


# 模型选择
k, num_epochs, lr, weight_decay, batch_size = 5, 100, 5, 0, 64
# 设置K折交叉验证的参数，训练轮数、学习率、正则化系数、batch size

train_l, valid_l = k_fold(k, train_features, train_labels, num_epochs, lr, weight_decay, batch_size)
# 执行K折交叉验证，返回训练集和验证集的平均误差

print('%d-fold validation: avg train rmse %f, avg valid rmse %f' % (k, train_l, valid_l))


# 打印K折交叉验证的训练和验证误差

# 定义函数
def train_and_pred(train_features, test_features, train_labels, test_data, num_epochs, lr, weight_decay, batch_size):
    net = get_net(train_features.shape[1])  # 创建模型
    train_ls, _ = train(net, train_features, train_labels, None, None, num_epochs, lr, weight_decay, batch_size)
    # 训练模型并获取训练误差
    d2l.semilogy(range(1, num_epochs + 1), train_ls, 'epochs', 'rmse', plot=True)
    # 绘制训练误差曲线
    print('train rmse %f' % train_ls[-1])  # 打印训练集的最终误差
    preds = net(test_features).detach().numpy()  # 预测测试集的房价，并转换为NumPy数组
    test_data['SalePrice'] = pd.Series(preds.reshape(1, -1)[0])
    # 将预测的房价添加到测试数据中
    submission = pd.concat([test_data['Id'], test_data['SalePrice']], axis=1)
    # 创建提交文件：包括Id和预测的房价
    submission.to_csv('./submission.csv', index=False)  # 将提交文件保存为CSV格式


train_and_pred(train_features, test_features, train_labels, test_data, num_epochs, lr, weight_decay, batch_size)
# 调用train_and_pred函数，训练模型并生成预测结果，保存到submission.csv文件中
